// Copyright 2024 Vadim Kharitonovich
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/ilyakaznacheev/cleanenv"
	"github.com/overflowingd/emlb"
	"github.com/rs/zerolog"
	zl "github.com/rs/zerolog/log"
	"notabug.org/overflowingd/lbs"
	"notabug.org/overflowingd/lbs/pkg/config"
	"notabug.org/overflowingd/lbs/pkg/log"
	"notabug.org/overflowingd/lbs/pkg/log/adapter"
	"notabug.org/overflowingd/lbs/pkg/protocol/packet"
	"notabug.org/overflowingd/lbs/pkg/server"
)

func StartExportHandlerState(ctx context.Context, handler *Handler, logger log.Logger) {
	go func() {
		ticker := time.NewTicker(time.Second * 15)

		for {
			select {
			case <-ctx.Done():
				return

			case <-ticker.C:
				retentions := uint64(handler.BalancerState.Retentions.Load())
				logger.Info("handler state", log.Any("balancer", map[string]any{
					"items_active":   handler.balancer.Len() - retentions,
					"items_retained": retentions,
				}))
			}
		}
	}()
}

type BalancerState struct {
	Retentions atomic.Int64
}

type Handler struct {
	BalancerState
	logger    log.Logger
	retention *Retention
	balancer  emlb.Balancer[[]byte]
}

func NewHandler(
	logger log.Logger,
	balancer emlb.Balancer[[]byte],
	retention *Retention,
) *Handler {
	return &Handler{
		logger:    logger,
		balancer:  balancer,
		retention: retention,
	}
}

func (r *Handler) Next(ctx context.Context, req *packet.Header[packet.RequestMessageCode]) *packet.Response {
	paylod, index, err := r.balancer.Next()
	if err != nil {
		return packet.NewResponseError(err)
	}

	return packet.NewResponseNext(index, paylod)
}

func (r *Handler) Retain(ctx context.Context, req *packet.Request[*packet.Retention]) *packet.Response {
	ok, err := r.balancer.Retain(req.Payload.Index())
	if err != nil {
		r.logger.Error("balancer retain", log.Err(err))
		return nil
	}

	if ok {
		r.BalancerState.Retentions.Add(1)
		r.logger.Debug("balancer retain", log.Uint64("index", req.Payload.Index()), log.Dur("duration", r.retention.Duration))

		go func(index uint64) {
			time.Sleep(r.retention.Duration)
			r.logger.Debug("balancer recover", log.Uint64("index", index))

			ok := r.balancer.Recover(index)
			if ok {
				r.BalancerState.Retentions.Add(-1)
			}
		}(req.Payload.Index())
	}

	return nil // client does not wait any response for retention
}

type Retention struct {
	Duration time.Duration `env:"RETENTION_DURATION" json:"duration"`
}

type Workers struct {
	Max int `env:"WORKERS_MAX" json:"max"`
	Min int `env:"WORKERS_MIN" json:"min"`
}

type Logger struct {
	Level string `env:"LOGGER_LEVEL" json:"level"`
}

type Options struct {
	Workers   `json:"workers"`
	Retention `json:"retention"`
	Logger    `json:"logger"`

	Listen         string `env:"LISTEN" json:"listen"`
	ConfigProvider string `env:"CONFIG_PROVIDER" json:"config_provider"`
	ConfigSource   string `env:"CONFIG_SOURCE" json:"config_source"`
}

func ProvideOptions() (*Options, error) {
	var (
		options Options
	)

	if err := cleanenv.ReadEnv(&options); err != nil {
		return nil, fmt.Errorf("cleanenv read env: %w", err)
	}

	if options.Listen == "" {
		flag.StringVar(&options.Listen, "listen", ":2222", "listen address in format <host>:<port>")
	}

	if options.ConfigProvider == "" {
		flag.StringVar(&options.ConfigProvider, "config-provider", "http-get", "method to use to fetch service config")
	}

	if options.ConfigSource == "" {
		flag.StringVar(&options.ConfigSource, "config-source", "", "URL where to fetch service config")
	}

	if options.Workers.Min < 1 {
		flag.IntVar(&options.Workers.Min, "workers-min", 5, "minimum number of workers concurrently handling packets")
	}

	if options.Workers.Max < 1 {
		flag.IntVar(&options.Workers.Max, "workers-max", 20, "maximum number of workers concurrently handling packets")
	}

	if options.Retention.Duration < 1 {
		flag.DurationVar(&options.Retention.Duration, "retention-duration", time.Minute*5, "period for which to retain items from balancing algorithm")
	}

	if options.Logger.Level == "" {
		flag.StringVar(&options.Logger.Level, "logger-level", zerolog.DebugLevel.String(), "logger verbosity level")
	}

	flag.Parse()

	if options.ConfigSource == "" {
		return nil, errors.New("flag config-source: required")
	}

	return &options, nil
}

func HandleTerm(cancel context.CancelFunc) {
	signals := make(chan os.Signal, 1)
	defer close(signals)
	signal.Notify(signals, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
	<-signals
	cancel()
}

func main() {
	zerolog.TimeFieldFormat = zerolog.TimeFormatUnix

	var (
		logger      = adapter.NewZerolog(zl.Level(zerolog.FatalLevel))
		ctx, cancel = context.WithCancel(context.Background())
	)

	go HandleTerm(cancel)

	options, err := ProvideOptions()
	if err != nil {
		logger.Fatal("provide options", log.Err(err))
	}

	level, err := zerolog.ParseLevel(options.Logger.Level)
	if err != nil {
		logger.Fatal("zerolog parse level", log.Err(err))
	}

	logger = adapter.NewZerolog(zl.Level(level))

	configSource, err := url.Parse(options.ConfigSource)
	if err != nil {
		logger.Fatal("parse config-source", log.Err(err))
	}

	config, err := config.Load(ctx, config.ProviderName(options.ConfigProvider), configSource)
	if err != nil {
		logger.Fatal("config load", log.Err(err))
	}

	logger.Info("config load", log.Str("status", "ok"))

	algorithm, err := emlb.NewRoundRobin(uint64(len(config.Items)))
	if err != nil {
		logger.Fatal("emlb new roundrobin", log.Err(err))
	}

	balancer, err := emlb.New(algorithm, config.Items)
	if err != nil {
		logger.Fatal("emlb new", log.Err(err))
	}

	handler := NewHandler(logger, balancer, &options.Retention)
	if logger.Enabled(log.InfoLevel) {
		StartExportHandlerState(ctx, handler, logger)
	}

	var (
		server = lbs.NewServer(
			handler,
			server.WithLogger(logger),
			server.WithMinWorkers(options.Workers.Min),
			server.WithMaxWorkers(options.Workers.Max),
		)
	)

	logger.Info(
		"server listen",
		log.Any("workers", options.Workers),
		log.Str("listen", options.Listen),
	)

	if err := server.Listen(ctx, options.Listen); err != nil {
		logger.Fatal("server listen", log.Err(err))
	}

	<-ctx.Done()
}
